#!/usr/bin/env python3
from fastmcp import Context
import asyncio
import time
from connection_manager import ConnectionManager
from typing import Dict, Optional
from logger import Logger
import re

support_elicitation = False


class ClientManager:
    def __init__(self, logger: Logger):
        self.clients: Dict[str, ConnectionManager] = {}
        self.lock = asyncio.Lock()
        self.scan_task = None
        self.scan_running = False
        self.logger = logger

    async def get_conn_manager(self, context: Context) -> [str, ConnectionManager]:
        """获取连接管理器"""
        # 获取客户端ID
        if context.session_id is None:
            return "Session ID 不存在", None
        conn_manager = await self.get_client_connection_manager(context.session_id)

        if not support_elicitation:
            if conn_manager is None:
                return "新创建连接，请使用工号登录，否则无法使用", conn_manager
        else:
            if conn_manager is None:
                for _ in range(0,3):
                    result = await context.elicit(
                        message="新创建连接，请输入工号，否则无法使用",
                        response_type=str
                    )
                    
                    if result.action == "accept":
                        user_id = result.data
                        # 验证工号格式：6位数字，前面可以补充0
                        if user_id is None or not re.match(r'^\d{6}$', user_id):
                            await context.log(f"工号格式不正确，必须是6位数字")
                            continue
                        await context.log(f"Hello {user_id}")
                        return "Success", await self.get_client_connection_manager_by_user_id(context.session_id, user_id)

        return "", conn_manager

    async def _create_client_internal(self, client_id: str, user_id: str) -> bool:
        """内部方法：创建新客户端（假设已持有锁）"""
        if client_id in self.clients:
            return False
        self.clients[client_id] = ConnectionManager(client_id, user_id)
        if not self.scan_task:
            self.scan_task = asyncio.create_task(self.start_scanner())
        return True

    async def create_client(self, client_id: str, user_id: str) -> bool:
        """创建一个新的客户端
        """
        async with self.lock:
            return await self._create_client_internal(client_id, user_id)

    async def get_client_connection_manager_by_user_id(self, new_client_id: str, user_id: str) -> Optional[ConnectionManager]:
        """根据用户ID查找客户端连接管理器
        """
        async with self.lock:
            for client_id, conn_manager in self.clients.items():
                if conn_manager.user_id == user_id:
                    del self.clients[client_id]
                    self.clients[new_client_id] = conn_manager
                    conn_manager.client_id = new_client_id
                    return conn_manager
            await self._create_client_internal(new_client_id, user_id)
        return self.clients[new_client_id]

    async def get_client_connection_manager(self, client_id: str) -> Optional[ConnectionManager]:
        """获取指定客户端的连接管理器
        参数:
        - client_id: 客户端标识符
        """
        async with self.lock:
            if client_id not in self.clients:
                return None
            return self.clients[client_id]

    async def _close_client_internal(self, client_id: str) -> int:
        """内部方法：关闭客户端（假设已持有锁）"""
        if client_id not in self.clients:
            return 0
        conn_manager = self.clients[client_id]
        count = await conn_manager.close_all_connections()
        # 删除客户端
        del self.clients[client_id]
        if len(self.clients) == 0:
            await self.stop_scanner()
        return count

    async def close_client(self, client_id: str) -> int:
        """关闭指定客户端的所有连接
        参数:
        - client_id: 客户端ID
        返回关闭的连接数量
        """
        async with self.lock:
            return await self._close_client_internal(client_id)

    async def close_all_clients(self) -> int:
        """关闭所有客户端的所有连接
        返回关闭的连接数量
        """
        # 获取客户端ID列表
        async with self.lock:
            client_ids = list(self.clients.keys())
        
        # 并行关闭所有客户端
        tasks = []
        for client_id in client_ids:
            tasks.append(self.close_client(client_id))
        
        # 等待所有任务完成
        results = await asyncio.gather(*tasks)
        return sum(results)
    
    async def start_scanner(self):
        """启动扫描任务，每小时扫描一次"""
        self.scan_running = True
        while self.scan_running:
            await self.scan_expired_clients()
            # 等待1小时
            for _ in range(360):  # 360个10秒
                if not self.scan_running:
                    break
                await asyncio.sleep(10)
    
    async def stop_scanner(self):
        """停止扫描任务"""
        self.scan_running = False
        if self.scan_task:
            # 不需要取消任务，让它自然结束
            self.scan_task = None
    
    async def scan_expired_clients(self):
        """扫描并关闭过期的客户端连接管理器"""
        async with self.lock:
            # 创建要关闭的客户端ID列表
            expired_client_ids = []
            for client_id, conn_manager in self.clients.items():
                if conn_manager.is_expired():
                    expired_client_ids.append(client_id)
        
        # 并行关闭所有过期的客户端
        if expired_client_ids:
            tasks = []
            for client_id in expired_client_ids:
                client = self.clients[client_id]
                self.logger.server_log(f"{client.user_id} 客户端，超时删除")
                tasks.append(self.close_client(client_id))
            
            if tasks:
                await asyncio.gather(*tasks)
                self.logger.server_log(f"已清理 {len(expired_client_ids)} 个过期的客户端连接管理器")

    async def list_clients(self) -> Dict[str, list]:
        """列出所有客户端及其连接
        返回格式: {client_id: [connection_info]}
        """
        # 获取客户端和连接管理器的副本
        async with self.lock:
            clients_copy = list(self.clients.items())
        
        # 并行获取每个客户端的连接信息
        result = {}
        tasks = {}
        for client_id, conn_manager in clients_copy:
            tasks[client_id] = asyncio.create_task(conn_manager.list_connections())
        
        # 等待所有任务完成并收集结果
        for client_id, task in tasks.items():
            result[client_id] = await task
        
        return result